stacked_prob_plot2 <- function (x_vals,
                                x_lab,
                                est,
                                offset = 0.04,
                                labels = NULL,
                                plot.new = TRUE) {
  if (length(dim(est)) == 3) {
    if (plot.new) {
      # plot.new()
      par(mar = c(0, 0, 0, 0),
          oma = c(5, 5, .1, .1))
    }
    plot(
      x_vals,
      x_vals,
      type = 'n',
      main = "",
      axes = F,
      ylab = "",
      xlab = "",
      ylim = c(0, 1),
      xlim = range(x_vals)
    )
    box()
    
    ## Predictions
    for (j in seq_len(dim(est)[2])) {
      if (j == 1) {
        polygon(
          c(x_vals, rev(x_vals)),
          c(rep(0, length(x_vals)),
            rev(est[1, 1,])),
          col = adjustcolor("gray10", alpha.f = .75),
          border = NA
        )
      } else {
        polygon(
          c(x_vals, rev(x_vals)),
          c(est[1, j - 1,],
            rev(est[1, j,])),
          col = adjustcolor(paste0("gray", j * 10), alpha.f = .75),
          border = NA
        )
      }
    }
    
    ## Confidence intervals
    for (j in seq_len(dim(est)[2] - 1)) {
      polygon(
        c(x_vals, rev(x_vals)),
        c(est[2, j, ],
          rev(est[3, j, ])),
        col = adjustcolor("white", alpha.f = 0.25),
        border = NA
      )
      lines(x_vals,
            est[1, j, ],
            col = adjustcolor("white", alpha.f = 0.75),
            lwd = 0.5)
    }
    rect(min(x_vals),
         0,
         max(x_vals),
         1,
         col = NA,
         border = "black")
    
    mid <- ceiling(0.5 * length(x_vals))
    x_mid <- x_vals[mid] + 
      ifelse(seq_len(dim(est)[2]) %% 2 == 0, -1, 1) * 
      offset
    y_mid <- c(mean(c(0, est[1, 1, mid])))
    for (j in 2:dim(est)[2]) 
      y_mid <- c(y_mid, mean(c(est[1, j - 1, mid], est[1, j, mid])))
    
    text(
      rep(x_mid, length(y_mid)),
      y_mid,
      labels = labels,
      cex = 0.75,
      col = "white"
    )
    axis(2, outer = TRUE)
    axis(1, outer = TRUE)
    title(xlab = x_lab, outer = TRUE)
    title(ylab = "Stacked probabilities", outer = TRUE)
  }
}